import os
import numpy as np
import nibabel as nib
from tqdm import tqdm


class SelectSlice:
    """
    Finds the first non-zero slice and takes the +4 slice.
    """
    def __init__(self, increment=4):
        self.increment = increment

    def __call__(self, segmentation):
        for i in range(segmentation.shape[2]):
            if np.any(segmentation[:, :, i, :]):
                target_slice = i + self.increment
                if target_slice < segmentation.shape[2]:
                    selected_segmentation = segmentation[:, :, target_slice, :]
                    selected_segmentation = np.transpose(selected_segmentation, axes=(2, 0, 1))
                    return selected_segmentation  # (T, H, W)
                break
        print(f"No valid slice found, or slice index {target_slice} is out of bounds.")
        return None

def find_ES_timestep(segmentation):
    lv_areas = np.zeros(segmentation.shape[0])
    rv_areas = np.zeros(segmentation.shape[0])
    
    for t in range(segmentation.shape[0]):
        lv_areas[t] = np.sum(segmentation[t, :, :] == 1)
        rv_areas[t] = np.sum(segmentation[t, :, :] == 3)
    
    # ES frame is when the sum of LV and RV areas is minimized
    ES_criteria = lv_areas + rv_areas
    ES_time_step = np.argmin(ES_criteria)
    
    return ES_time_step

def process_single_folder(folder_path):
    seg_path = os.path.join(folder_path, 'seg_sa.nii.gz')

    if os.path.exists(seg_path):
        seg_vol = nib.load(seg_path).get_fdata()
        selector = SelectSlice(increment=4)
        selected_segmentation = selector(seg_vol)

        if selected_segmentation is not None:
            ES_time_step = find_ES_timestep(selected_segmentation)
            return ES_time_step
    else:
        print(f"No seg or folder for {folder_path}")
        return -1

def process_folders(root_dir, ids):
    folder_paths = [os.path.join(root_dir, str(i)) for i in ids]
    es_time_steps = []

    for folder_path in tqdm(folder_paths):
        es_time_step = process_single_folder(folder_path)
        es_time_steps.append(es_time_step)
    
    return np.array(es_time_steps, dtype=np.int32)

if __name__ == '__main__':
    root_dir = '/vol/aimspace/projects/ukbb/data/cardiac/cardiac_segmentations/subjects/'
    phases = ['train', 'val', 'test']
    
    for phase in phases:
        ids_path = f'saved_tensors/ids/{phase}_patient_ids.npy'
        ids = np.load(ids_path)
        es_time_steps_array = process_folders(root_dir, ids)
        
        save_path = f'saved_tensors/multimodal/imaging_ES_frames/imaging_{phase}_es_timestep.npy'
        np.save(save_path, es_time_steps_array)
